import numpy as np


def randMatrix(PMD_para):
    PSP_theta = PMD_para["PSP_theta"]
    PSP_phi = PMD_para["PSP_phi"]

    # 泡利矩阵
    Pauli_1 = np.eye(2)
    Pauli_2 = np.array([[0, 1], [1, 0]])
    Pauli_3 = np.array([[0, 1], [-1, 0]])

    PSP_mat_theta = np.cos(PSP_theta) * Pauli_1 - np.sin(PSP_theta) * Pauli_3
    PSP_mat_phi = np.cos(PSP_phi) * Pauli_1 + 1j * np.sin(PSP_phi) * Pauli_2

    # PSP的旋转
    Matrix = np.dot(PSP_mat_theta, PSP_mat_phi)

    return Matrix


def Trunk_Matrix(Signal, Fiber_para, PMD_para):
    # 随机波片的光谱分解
    # 生成随机的Jones矩阵，使相邻trunk之间发生随机的旋转，模拟随机偏振模色散
    # Jones矩阵是一个2*2的矩阵，共生成trunk个不同的矩阵，每个矩阵对应一个特征值
    # Matrix对应2*2矩阵
    # Value_num对应特征值

    Npol = Signal["polarization"]

    L = Fiber_para["L"]
    # Beatlength = PMD_para["beatlength"] #极化间的比特长度
    # 相干长度
    L_corr = PMD_para["L_corr"]
    # 单个Span的总差分群时延
    PMD_DGD = PMD_para["PMD_DGD"]
    # 拍长[m]
    beatlength = PMD_para['beatlength']
    # 单位长度的差分群时延[ns / km]
    DGD_unit = PMD_DGD * 1e12 / np.sqrt(L_corr * 1e3) * np.sqrt(3 * np.pi / 8) / np.sqrt(1000)
    # trunk_length = PMD_para["trunk_length"]
    # trunk_num = PMD_para["trunk_num"]
    # Maxwell distribution
    q = np.sqrt(np.pi ** 3 / 2) / beatlength * np.random.randn(1, 3)
    # %deltabeta0[1 / m]
    w = np.sqrt(np.sum(q ** 2))
    # 特征值
    eigenvalue = np.arange(0.1, 0.3, 0.1)
    eigenvalue[0] = w / 2
    eigenvalue[1] = -w / 2

    PMD_para['DGD_unit'] = DGD_unit

    PMD_para["PSP_theta"] = np.random.rand(1) * 2 * np.pi - np.pi
    PMD_para["PSP_phi"] = 0.5 * np.arcsin(np.random.rand(1) * 2 - 1)

    if Npol == 1:
        Matrix = np.eye(2)  # 将矩阵赋值为2*2的单位阵，表示波片的特征向量为正交基
    else:
        # 生成PMD参数
        Matrix = randMatrix(PMD_para)  # 随机生成酉矩阵

    return Matrix, eigenvalue, PMD_para
